from pypcd import pypcd
from point_to_bev import *
import os
import json
from scipy.spatial.transform import Rotation as R

def rotation_matrix_x(theta):
    """绕x轴的旋转矩阵"""
    return np.array([
        [1, 0, 0],
        [0, np.cos(theta), -np.sin(theta)],
        [0, np.sin(theta), np.cos(theta)]
    ])

def rotation_matrix_y(theta):
    """绕y轴的旋转矩阵"""
    return np.array([
        [np.cos(theta), 0, np.sin(theta)],
        [0, 1, 0],
        [-np.sin(theta), 0, np.cos(theta)]
    ])

def rotation_matrix_z(theta):
    """绕z轴的旋转矩阵"""
    return np.array([
        [np.cos(theta), -np.sin(theta), 0],
        [np.sin(theta), np.cos(theta), 0],
        [0, 0, 1]
    ])

def calculate_vertices(x, y, z, l, w, h, roll,pitch,yaw):
    # 长方体的半尺寸
    half_l = l / 2
    half_w = w / 2
    half_h = h / 2

    # 未旋转时的顶点坐标
    vertices = np.array([
        [x - half_l, y - half_w, z + half_h],
        [x + half_l, y - half_w, z + half_h],
        [x + half_l, y + half_w, z + half_h],
        [x - half_l, y + half_w, z + half_h],
        [x - half_l, y - half_w, z - half_h],
        [x + half_l, y - half_w, z - half_h],
        [x + half_l, y + half_w, z - half_h],
        [x - half_l, y + half_w, z - half_h]
    ])

    # 绕z轴旋转的变换矩阵
    rx = rotation_matrix_x(roll)
    ry = rotation_matrix_y(pitch)
    rz = rotation_matrix_z(yaw)

    # 旋转顶点坐标
    # rotated_vertices = np.dot(vertices - np.array([x, y, z]), rx) + np.array([x, y, z])
    # rotated_vertices = np.dot(vertices - np.array([x, y, z]), ry) + np.array([x, y, z])
    rotated_vertices = np.dot(vertices - np.array([x, y, z]), rz) + np.array([x, y, z])

    return rotated_vertices


def read_pcd(lidar_path):
    xyzi=np.load(lidar_path)[:,:4]
    bev_image=point_to_bev(xyzi)
    return bev_image

if __name__=="__main__":
    root_path="/data2/xd/CenterPoint/dataset/rcs_0422/samples/20240815/test0815"

    lidar_list=sorted(os.listdir(root_path+'/lidar'))
    label_path=os.path.join(root_path,"test.json")
    labels=json.load(open(label_path,"r"))

    for i in range(len(lidar_list)):
        obj=lidar_list[i]
        lidar_path=os.path.join(root_path,"lidar",obj)

        bev_image=read_pcd(lidar_path)



        for label in labels['items'][i]['annotations']:
            x,y,z=label['position'][0],label['position'][1],label['position'][2]
            roll,pitch,yaw=label['rotation'][0],label['rotation'][1],label['rotation'][2]
            l,w,h=label['scale'][0],label['scale'][1],label['scale'][2]
            
            points=calculate_vertices(x, y, z, l, w, h, roll,pitch,np.pi-yaw)
            
            points[:,0]=1080/2+points[:,0]*1080/60
            points[:,1]=1920/2-points[:,1]*1920/60
            points=points.astype(np.int32)

            cv2.line(bev_image,(points[0][0],points[0][1]) , (points[1][0],points[1][1]) ,(0,255,255),1)
            cv2.line(bev_image,(points[1][0],points[1][1]) , (points[2][0],points[2][1]) ,(0,0,255),1)
            cv2.line(bev_image,(points[2][0],points[2][1]) , (points[3][0],points[3][1]) ,(0,0,255),1)
            cv2.line(bev_image,(points[3][0],points[3][1]) , (points[0][0],points[0][1]) ,(0,0,255),1)
        cv2.imwrite("bevtest.jpg",bev_image)



